import torch
from torch import nn, tensor
from garbageDataset import garbageDataset
from torch.utils.data import DataLoader, Dataset
import numpy as np
import pandas as pd
import cv2 as cv
import os
import torch.onnx

class NervAlexNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_channels=3, out_channels=96, kernel_size=11, stride=4, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(in_channels=96, out_channels=256, kernel_size=5, padding=2),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(in_channels=256, out_channels=384, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.Conv2d(in_channels=384, out_channels=384, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.Conv2d(in_channels=384, out_channels=256, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=3, stride=2)
        )
        self.flat = lambda x: x.view(x.size(0), -1)
        self.full = nn.Sequential(
            nn.Linear(in_features=6400, out_features=4096),
            nn.ReLU(),
            nn.Dropout(0.4),
            nn.Linear(in_features=4096, out_features=4096),
            nn.ReLU(),
            nn.Dropout(0.4),
            nn.Linear(in_features=4096, out_features=160)
        )
        self.gpus = None

    def forward(self, x):
        x = self.conv(x)
        x = self.flat(x)
        x = self.full(x)
        return x

    def Train(self, epochs, Lr, Bs, saveOpt=False, savePath=None, DstrOpt=False):
        citerition = nn.CrossEntropyLoss()
        optimzer = torch.optim.Adam(self.parameters(), lr=Lr)
        data = DataLoader(garbageDataset(folderPath="./trainData/", size=224), batch_size=Bs, num_workers=4, shuffle=True)
        if DstrOpt:
            self.gpus = self.Gpus()
            nn.DataParallel(self, device_ids=self.gpus)

        for epoch in range(epochs):
            correct, all, Loss, bestacc, times = 0, 0, 0.0, 0.0, 0

            for idx, (x, y) in enumerate(data):
                if DstrOpt:
                    x, y = x.to(self.gpus[0]), y.to(self.gpus[0])

                if idx % 10 < 7:
                    self.train()
                    predict = self(x)
                    loss = citerition(predict, y)
                    loss.backward()
                    optimzer.step()
                    optimzer.zero_grad()

                else:
                    all += Bs
                    self.eval()
                    predict = self(x)
                    Loss += citerition(predict, y)
                    _, predicted = torch.max(predict.data, 1)
                    correct += (predicted == y).sum()

                    acc = correct / (1.0 * all)
                    avg_loss = Loss / (1.0 * all)
                    print("\r>>epoch:{}/{}, acc:{:.2%}, avg_loss:{:.4f}.".format(epoch + 1, epochs, acc, avg_loss), end="")

            print()
            if saveOpt:
                if acc > 1.02 * bestAcc:
                    bestacc = acc
                    self.Save(savePath)

    def Gpus(self):
        devices = [torch.device(f'cuda:{i}') for i in range(torch.cuda.device_count())]
        return devices if devices else [torch.device('cpu')]

    def Save(self, filePath):
        torch.save(self.state_dict(), filePath)

    def Load(self, filePath):
        self.load_state_dict(torch.load(filePath, map_location=torch.device('cpu')))

if __name__ == "__main__":
    myNet = NervAlexNet()
    test = torch.rand(1, 3, 224, 224, dtype=torch.float)
    #myNet.Train(20, 0.0001, 16, DstrOpt=True)
    myNet.Load("./alex2")
    torch.onnx.export(model=myNet, args=test, f="./now.proto");
    
